# !/usr/bin/env python
# -*- coding: utf-8 -*-

# 请确保你的 DataFrame 包含一个名为 Label 的列, 这是 Species(品种) 的列。
# 如果 label 是字符串类型的分类特征, 你可能需要使用 StringIndexer 进行索引。
'''
任务类型: 随机森林回归主要用于回归任务。在回归任务中, 算法试图预测一个连续的数值输出, 而不是一个离散的类别。
输出: 随机森林回归的输出是一个连续的数值, 表示输入数据的预测结果。
算法原理: 随机森林回归同样基于决策树, 但在回归任务中, 每个决策树的输出是一个实数值。最终的预测结果是多个决策树输出的平均值或加权平均值。

字段说明: SepalLength(花萼长度), SepalWidth(花萼宽度), PetalLength(花瓣长度), PetalWidth(花瓣宽度), Species(品种).

品种说明: Setosa(山鸢尾), Versicolor(变色鸢尾), Virginical(维吉尼亚鸢尾).
'''
'''
跨平台生成模型，java调用预测

https://zhuanlan.zhihu.com/p/435694136
'''
from pyspark.ml import Pipeline, PipelineModel
from pyspark.sql import Row, SparkSession
from pyspark.sql.functions import col
from pyspark.sql.types import StringType, DoubleType
from pyspark.ml.feature import StringIndexer, VectorAssembler
from pyspark.ml.regression import RandomForestRegressor
from pyspark2pmml import PMMLBuilder

if __name__ == '__main__':
    spark = SparkSession.builder.master("local[*]").config("spark.jars.packages",
                                                           "org.jpmml:pmml-sparkml:2.4.1").appName(
        "model_learn").getOrCreate()

    # 加载数据
    df = spark.read.option("header", "true").option("encoding", "utf-8").csv("../data/iris.csv")
    df.show()

    # 转换 Spark 中 DateFrame 数据类型。
    df = df.withColumn("SepalLength", col("SepalLength").cast(DoubleType()))
    df = df.withColumn("SepalWidth", col("SepalWidth").cast(DoubleType()))
    df = df.withColumn("PetalLength", col("PetalLength").cast(DoubleType()))
    df = df.withColumn("PetalWidth", col("PetalWidth").cast(DoubleType()))
    df = df.withColumn("Species", col("Species").cast(StringType()))

    # 输出 Spark 中 DataFrame 字段和数据类型。
    print("Changed Spark DataFrame Data Type:")
    df.printSchema()

    # 使用 StringIndexer 转换 Species 列。
    indexer = StringIndexer(inputCol="Species", outputCol="SpeciesIdx")
    # 拟合并转换数据。
    indexDF = indexer.fit(df).transform(df)

    # 输出 StringIndexer 的转换效果。
    print("The Effect of StringIndexer:")
    indexDF.show()

    # 将 SpeciesIdx 列复制为Label列
    dataset = indexDF.withColumn("Label", col("SpeciesIdx"))

    # VectorAssembler 将多个特征合并为一个特征向量。
    required_features: list = ["SepalLength", "SepalWidth", "PetalLength", "PetalWidth"]
    assembler = VectorAssembler(inputCols=required_features, outputCol="Features")

    model_path = "../out/house_price.pmml"
    # 创建 Pipeline (可选): 将特征向量转换和随机森林回归模型组合在一起
    # 注意: 如果要使用 Pipeline, 则在创建 特征向量(Features)列 的时候不需要拟合数据, 否则会报 "Output column Features already exists." 的错误。
    # 创建 随机森林回归(RandomForestRegressor)。
    # 将数据集划分为 "训练集" 和 "测试集"。
    (trainingData, testData) = dataset.randomSplit([0.8, 0.2], seed=42)
    rf = RandomForestRegressor(featuresCol="Features", labelCol="Label")
    pipeline = Pipeline(stages=[assembler, rf])
    # 训练模型 (Pipeline 模式)。
    model = pipeline.fit(trainingData)
    pmmlBuilder = PMMLBuilder(spark.sparkContext, trainingData, model).putOption(rf, "compact", True)
    pmmlBuilder.buildFile(model_path)


    spark.stop()
